import os
import pathlib

from dotenv import load_dotenv
from sqlalchemy.ext.asyncio import AsyncEngine
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker, declarative_base

# 加载环境变量
load_dotenv()

# 获取数据库类型和连接信息
DB_TYPE = os.getenv("DB_TYPE", "sqlite").lower()  # 默认使用SQLite

# 数据库连接URL
if DB_TYPE == "mysql":
    # MySQL连接信息
    DB_HOST = os.getenv("DB_HOST", "localhost")
    DB_PORT = os.getenv("DB_PORT", "3306")
    DB_USER = os.getenv("DB_USER", "root")
    DB_PASSWORD = os.getenv("DB_PASSWORD", "123456")
    DB_NAME = os.getenv("DB_NAME_MYSQL", "test_db")

    # 创建MySQL异步数据库连接URL
    DATABASE_URL = f"mysql+aiomysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
else:
    # SQLite连接信息
    # 确保数据目录存在
    DB_DIR = os.getenv("DB_DIR", "./db")
    pathlib.Path(DB_DIR).mkdir(parents=True, exist_ok=True)

    DB_NAME = os.getenv("DB_NAME_SQLITE", "test_db.db3")
    DB_PATH = os.path.join(DB_DIR, DB_NAME)

    # 创建SQLite异步数据库连接URL
    DATABASE_URL = f"sqlite+aiosqlite:///{DB_PATH}"

# 创建异步引擎
engine: AsyncEngine = create_async_engine(
    DATABASE_URL,
    echo=True,  # 设置为True以查看SQL语句
    pool_pre_ping=True,  # 连接池健康检查
    pool_recycle=3600,  # 连接回收时间
)

# 创建异步会话工厂
AsyncSessionLocal = sessionmaker(
    engine, class_=AsyncSession, expire_on_commit=False
)

# 创建Base类，用于创建模型类
Base = declarative_base()


# 异步上下文管理器，用于获取数据库会话
async def get_db():
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception as e:
            await session.rollback()
            raise e


# 初始化数据库
async def init_db():
    async with engine.begin() as conn:
        # 创建所有表
        await conn.run_sync(Base.metadata.create_all)


# 关闭数据库连接
async def close_db():
    await engine.dispose()
